from deepface import DeepFace
import time
import tensorflow as tf
import cv2

# --------------------------
# 1. GPU配置优化（关键加速点）
# --------------------------
# 启用GPU内存动态增长（避免一次性占用全部显存，提升资源利用率）
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
        print("GPU配置成功：", gpus)
    except RuntimeError as e:
        print(e)

# 启用混合精度推理（利用RTX 5070的FP16算力，加速20-30%）
tf.keras.mixed_precision.set_global_policy('mixed_float16')


# --------------------------
# 2. 图像预处理优化（减少计算量）
# --------------------------
def preprocess_image(img_path, target_size=(150, 150)):
    """提前预处理图像，统一尺寸并减少噪声"""
    img = cv2.imread(img_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # 转换为RGB格式（DeepFace默认输入）
    img = cv2.resize(img, target_size)  # 缩小尺寸（根据模型需求调整，平衡速度与精度）
    img = img / 255.0  # 归一化
    return img


# --------------------------
# 3. 模型与检测器选择（平衡速度与精度）
# --------------------------
# 模型选择：Facenet512比OpenFace推理更快，且精度相当；MobileFaceNet更轻量（适合极致速度）
# 检测器选择：retinaface比mtcnn快30%，且检测精度接近
MODEL_NAME = "Facenet512"  # 备选："MobileFaceNet"（更快）、"OpenFace"（原模型）
DETECTOR_BACKEND = "ssd"  # 备选："ssd"（最快但精度略低）、"mtcnn"（原检测器）


# --------------------------
# 4. 核心验证逻辑（复用模型避免重复加载）
# --------------------------
def verify_faces(img1, img2):
    # 提前加载模型（DeepFace默认每次调用重新加载，预加载可节省20-30%时间）
    # 注意：首次运行仍会下载模型（约100-300MB），后续复用
    start_time = time.time()

    # 预处理图像（减少模型输入尺寸，加速特征提取）
    # img1 = preprocess_image(img1_path)
    # img2 = preprocess_image(img2_path)

    # 调用验证函数（复用预处理后的图像）
    result = DeepFace.verify(
        img1_path=img1,  # 直接传入预处理后的数组，避免重复读取
        img2_path=img2,
        model_name=MODEL_NAME,
        detector_backend=DETECTOR_BACKEND,
        enforce_detection=False  # 若图像中无脸，返回False而非报错（根据需求调整）
    )

    end_time = time.time()
    return result, end_time - start_time


# --------------------------
# 执行与输出
# --------------------------
if __name__ == "__main__":
    # 图像路径
    img1_path = "person1_1.jpg"
    img2_path = "person1_2.jpg"  # 可替换为不同人图像测试

    # 执行验证
    result, search_time = verify_faces(img1_path, img2_path)

    # 输出结果
    print(f"\n搜索耗时: {search_time:.2f} 秒")
    print("是否匹配：", result["verified"])
    print("匹配置信度：", result["distance"])
    print("模型阈值：", result["threshold"])